package common

import (
	"bytes"
	"errors"
	"io"
	"log"
	"net"
	"regexp"
	"shell/common/packages"
	"shell/global"
	"shell/utils"

	"github.com/google/uuid"
)

/*
	socks5代理相关代码
*/

const (
	socks5Version  byte = 0x05 // socks5版本号
	authNone       byte = 0x00 // 无认证的代号
	authUserPasswd byte = 0x02 // 用户名密码认证的代号
	cmdConnect     byte = 0x01
	rsv            byte = 0x00 // 保留字段
)

var (
	ErrInvalidSocks5Format   = errors.New("error invalid socks5 format")
	ErrInvalidSocks5Version  = errors.New("error invalid socks5 version")
	ErrUnsupportedAuthMethod = errors.New("error unsupported auth method")
	ErrUnsupportedCmd        = errors.New("error unsupported command")
	ErrUnsupportedAddrType   = errors.New("error unsupported address type")

	// 回复客户端使用无认证的包
	RespNoAuth = []byte{socks5Version, authNone}

	// 回复客户端Socks5代理建立成功
	RespProxyOk = []byte{socks5Version, 0, rsv, byte(packages.ATypeIPv4), 0, 0, 0, 0, 0, 0}

	// regIPv6 匹配IPv6地址
	regIPv6 = regexp.MustCompile(".*:+.*")
)

// Socks5 处理Socks5客户端请求、生成请求结构体，如果error不为nil，就会关闭网络连接
// 注意，如果err为nil，这里也没有回复Socks5客户端
func Socks5(conn io.ReadWriteCloser) (*packages.Socks5RedirReq, error) {
	buf := make([]byte, 256)
	closeConn := func() { _ = conn.Close() }

	// 读取客户端的认证请求
	_, err := io.ReadFull(conn, buf[:2])
	if err != nil {
		closeConn()
		return nil, err
	}
	if buf[0] != socks5Version {
		closeConn()
		return nil, ErrInvalidSocks5Version
	}
	// 支持认证方法的数量
	nmethods := buf[1]
	_, err = io.ReadFull(conn, buf[:nmethods])
	if err != nil {
		closeConn()
		return nil, err
	}
	// 是否有无认证方法
	hasNoAuth := false
	for _, v := range buf[:nmethods] {
		if v == authNone {
			hasNoAuth = true
			break
		}
	}
	if !hasNoAuth {
		closeConn()
		return nil, ErrUnsupportedAuthMethod
	}

	// 回复选用无认证方法
	_, err = conn.Write(RespNoAuth)
	if err != nil {
		closeConn()
		return nil, err
	}

	_, err = io.ReadFull(conn, buf[:4])
	if err != nil {
		closeConn()
		return nil, err
	}

	if buf[0] != socks5Version {
		closeConn()
		return nil, ErrInvalidSocks5Version
	}
	if buf[1] != cmdConnect {
		closeConn()
		return nil, ErrUnsupportedCmd
	}
	if buf[2] != rsv {
		closeConn()
		return nil, ErrInvalidSocks5Format
	}

	result := new(packages.Socks5RedirReq)
	result.CmdUUID = uuid.New()
	// 解析目标地址类型，IPv4、IPv6或者域名
	result.TargetAType = packages.AType(buf[3])
	switch result.TargetAType {
	case packages.ATypeIPv4:
		_, err := io.ReadFull(conn, buf[:4])
		if err != nil {
			closeConn()
			return nil, err
		}
		result.TargetAddr = bytes.Clone(buf[:4])
	case packages.ATypeDomain:
		_, err = io.ReadFull(conn, buf[:1])
		if err != nil {
			closeConn()
			return nil, err
		}
		domainLen := int(buf[0])
		_, err = io.ReadFull(conn, buf[:domainLen])
		if err != nil {
			closeConn()
			return nil, err
		}
		targetAddr := string(buf[:domainLen])
		ip := net.ParseIP(targetAddr)
		if ip != nil {
			// 符合IP地址格式
			if regIPv6.Match([]byte(ip.String())) {
				// IPv6地址
				result.TargetAType = packages.ATypeIPv6
				result.TargetAddr = ip.To16()
			} else {
				// IPv4地址
				result.TargetAType = packages.ATypeIPv4
				result.TargetAddr = ip.To4()
			}
		} else {
			result.TargetAddr = []byte(targetAddr)
		}
	case packages.ATypeIPv6:
		_, err = io.ReadFull(conn, buf[:16])
		if err != nil {
			closeConn()
			return nil, err
		}
		result.TargetAddr = bytes.Clone(buf[:16])
	default:
		closeConn()
		return nil, ErrUnsupportedAddrType
	}

	// 读取目标端口，端口号是按照大端序传输的
	_, err = io.ReadFull(conn, buf[:2])
	if err != nil {
		closeConn()
		return nil, err
	}
	result.TargetPort, _ = utils.BytesToU16(buf[:2])
	result.Role = packages.RoleClient
	return result, nil
}

// 向客户端(from)回复Socks5连接成功，并开始转发
func RespSocksAndRedir(from, to io.ReadWriteCloser, tw *TimeoutWatcher) {
	if tw == nil {
		log.Fatalf("timeout watcher is nil %s", global.CodePosition(0, 0))
	}
	_, err := from.Write(RespProxyOk)
	if err != nil {
		log.Printf("error write socks5 resp: %v", err)
		return
	}
	NewTimeoutRWCPairWarp(from, to, global.DefaultIdleTime, tw).RedirSyncNoKey()
}
